#!/usr/bin/env python
#
# Stoker: an LDAP-aware utility for cluster/networked system management
#
# Author:   Mike Murphy (mamurph@cs.clemson.edu)
#           Michael Fenn (mfenn@cs.clemson.edu)
# Revision: 16 September 2008
#
#   Copyright 2007-2008 Clemson University
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
#
# This material is based upon work supported under a National Science
# Foundation Graduate Research Fellowship. Any opinions, findings,
# conclusions or recommendations expressed in this publication are those of
# the author(s) and do not necessarily reflect the views of the National
# Science Foundation.

import sys, os, os.path, getopt, signal, threading

# DELAY import of stoker_ldap until the user requests it... this way, a missing
# python-ldap installation does not break Stoker
import stoker_resolvers

from stoker_defines import CFG_USE_LDAP, DFLT_USE_LDAP, \
     EXIT_ERROR_LDAP, CFG_RESOLVE_HOSTNAME, DFLT_RESOLVE_HOSTNAME, \
     SHORT_OPTIONS, LONG_OPTIONS, ARGUMENTS, MIN_ARGS, MAX_ARGS, \
     EXIT_ERROR_SYNTAX, RESOLVE_TO_HOSTNAME, RESOLVE_TO_IP, RESOLVER_RULE, \
     RESOLVE_TO_DEFAULT, EXIT_ERROR_TARGET, CMD_CLASS, HOSTNAME_KEY, \
     IPV4_KEY, MAC_KEY, RESOLVE_TO_MAC
#

from stoker_messages import MSG_LDAP_MISSING, MSG_LDAP_ERROR, MSG_MIN_ARGS, \
     MSG_MAX_ARGS, MSG_UNKNOWN_COMMAND, MSG_UNKNOWN_TARGET, \
     MSG_UNRESOLVED_HOSTNAME, MSG_UNRESOLVED_MAC
#

from stoker_base import pr_error, pr_summary, make_report, write_file, \
                        pr_warning
#

from stoker_commands import STOKER_COMMAND_MAP

# OutputHandler: A limited monitor that receives exit codes, standard output,
# and standard error from each execution thread. If summary printing is
# enabled, a summary is printed to stderr at each receipt. Applying str()
# to an object of this class will return a detailed output report, with
# child stdout and stderr results (if present), sorted by machine name.
class OutputHandler(object):
   def __init__(self, summary = True, stdout_path = None, stderr_path = None):
      self.summary = summary
      self.sema = threading.BoundedSemaphore(1)
      self.global_code = 0
      self.results = {}
      self.stdout_path = stdout_path
      self.stderr_path = stderr_path
   #
   def put_result(self, who, code, out, err):
      self.sema.acquire()
      self.results[who] = (code, out.strip(), err.strip())
      if code != 0 and self.global_code == 0:
         self.global_code = code
      if self.summary:
         pr_summary(who, code)
      if self.stdout_path:
         write_file(who, self.stdout_path, out.strip())
      if self.stderr_path:
         write_file(who, self.stderr_path, err.strip())
      self.sema.release()
   #
   def get_result(self, who):
      self.sema.acquire()
      result = (-1, '', '')
      if who in self.results:
         result = self.results[who]
      self.sema.release()
      return result
   #
   def __repr__(self):
      self.sema.acquire()
      report = ''
      keys = self.results.keys()
      keys.sort()
      for key in keys:
         report += make_report(key, *self.results[key])
      self.sema.release()
      return report
   #
   def get_global_code(self):
      self.sema.acquire()
      code = self.global_code
      self.sema.release()
      return code
   #
   def get_raw_results(self):
      self.sema.acquire()
      output = ''
      error = ''
      keys = self.results.keys()
      keys.sort()
      for key in keys:
         output += self.results[key][1] + '\n'
         error += self.results[key][2] + '\n'
      self.sema.release()
      return (output.strip(), error.strip())
#


# NEW in 0.31, execution now happens on a one thread per machine
# basis.
class ExecutionThread(threading.Thread):
   def __init__(self, config, output, cmdobj, who, sema, stdin_path=[]):
      threading.Thread.__init__(self)
      self.config = config
      self.output = output
      self.cmdobj = cmdobj
      self.who = who
      self.sema = sema
      self.stdin_path = stdin_path
   #
   def run(self):
      stdin_data = ''
      proceed = True
      code = 0
      err = ''
      out = ''
      if len(self.stdin_path) > 0:
         for item in self.stdin_path:
            try:
               fname = item
               if os.path.isdir(item):
                  fname = os.path.join(item, self.who)
               if os.path.isfile(fname):
                  fh = open(fname, 'r')
                  stdin_data = fh.read()
                  fh.close()
                  break
            except Exception, message:
               err = message
               proceed = False
               break
      if proceed:
         code, out, err = self.cmdobj.execute(stdin_data)
      else:
         code = 2
      self.output.put_result(self.who, code, out, err)
      # Signal thread complete
      self.sema.release()
   #
#


# Main Stoker class
class Stoker(object):
   def __init__(self, config, max_threads, output):
      self.config = config
      self.max_threads = max_threads
      self.output = output
      self.resolvers = stoker_resolvers.get_resolvers(self.config)
      self.resolve_to_ip = \
            config.getParam(CFG_RESOLVE_HOSTNAME, DFLT_RESOLVE_HOSTNAME)
   #
   # Split options and arguments, returning a tuple (opts, args)
   def split_opts_args(self, command, optargs):
      short_opts = STOKER_COMMAND_MAP[command][SHORT_OPTIONS]
      long_opts = STOKER_COMMAND_MAP[command][LONG_OPTIONS]
      arg_list = STOKER_COMMAND_MAP[command][ARGUMENTS]
      min_args = STOKER_COMMAND_MAP[command][MIN_ARGS]
      max_args = STOKER_COMMAND_MAP[command][MAX_ARGS]
      #
      try:
         opts, args = getopt.getopt(optargs, short_opts, long_opts)
      except getopt.GetoptError, message:
         pr_error(message, EXIT_ERROR_SYNTAX)
      #
      if len(args) < min_args:
         pr_error(MSG_MIN_ARGS % (min_args), EXIT_ERROR_SYNTAX)
      elif max_args >= 0 and len(args) > max_args:
         pr_error(MSG_MAX_AGRS % (max_args), EXIT_ERROR_SYNTAX)
      #
      return (opts, args)
   #
   # Return a list of resolved hostnames or IP addresses, with duplicates
   # removed. Also merge LDAP and local targets.
   #
   # Return format is (resolved, unresolved)
   # where each item in resolved is (resolved_target, all_info)
   def resolve_targets(self, target_list, resolver_list, resolution):
      resolved = []
      unresolved = target_list[:]
      #
      # NOTE: order matters here! If the user has limited the resolvers
      # ("sources" in user-speak), the resolution should be done in order
      # of the resolvers.
      for resolver in resolver_list:
         res, unres = resolver.resolveTargets(unresolved)
         unresolved = unres
         ok = True
         
         for item in res:
            resname = item[HOSTNAME_KEY]
            if resolution == RESOLVE_TO_IP:
               if IPV4_KEY in item:
                  resname = item[IPV4_KEY]
               else:
                  pr_warning(MSG_UNRESOLVED_HOSTNAME % (resname))
                  ok = False
            elif resolution == RESOLVE_TO_MAC:
               if MAC_KEY in item:
                  resname = item[MAC_KEY]
               else:
                  pr_warning(MSG_UNRESOLVED_MAC % (resname))
                  ok = False
            if ok and (resname, item) not in resolved:
               resolved.append((resname, item))
      #
      return resolved, unresolved
   #

   def execute(self, cmdlist, all_machines = False, stdin_path = []):
      command = ''
      target_list = []
      optargs = ''
      resolution = RESOLVE_TO_HOSTNAME
      if self.resolve_to_ip:
         resolution = RESOLVE_TO_IP
      if not all_machines:
         command = cmdlist[1]
         targets = cmdlist[0]
         optargs = cmdlist[2:]
         target_list = targets.split(',')
      else:
         command = cmdlist[0]
         optargs = cmdlist[1:]
         target_list.extend(stoker_resolvers.get_all_targets(self.resolvers))
      #
      if command not in STOKER_COMMAND_MAP:
         pr_error(MSG_UNKNOWN_COMMAND % (command), EXIT_ERROR_SYNTAX)
      elif STOKER_COMMAND_MAP[command][RESOLVER_RULE] != RESOLVE_TO_DEFAULT:
         resolution = STOKER_COMMAND_MAP[command][RESOLVER_RULE]
      #
      machines, unresolved = \
            self.resolve_targets(target_list, self.resolvers, resolution)
      if len(unresolved) > 0:
         errlist = []
         for target in unresolved:
            errlist.append(MSG_UNKNOWN_TARGET % (target))
         pr_error(errlist, EXIT_ERROR_TARGET)
      #

      opts, args = self.split_opts_args(command, optargs)
      retcode = 0
      thread_list = []

      # Prepare a bounding semaphore to limit the number of threads that
      # may start. Start threads as the semaphore permits.
      sema = threading.BoundedSemaphore(self.max_threads)

      # Prepare a thread for each machine, then start as many threads as are
      # permitted to run at once. Start new threads as threads finish, if any
      # remain.
      for machine, extra in machines:
         cmd_class = STOKER_COMMAND_MAP[command][CMD_CLASS]
         hostname = extra[HOSTNAME_KEY]
         cmdobj = cmd_class(self.config, opts, args, machine, extra)
         machine_thread = ExecutionThread(self.config, self.output, cmdobj, \
                                          hostname, sema, stdin_path)
         thread_list.append(machine_thread)
      #
      for thread in thread_list:
         try:
            sema.acquire()
         except Exception, data:
            print data
            sys.exit(EXIT_ERROR_SYNTAX)
         retcode = thread.start()
      #
      # Now wait for all threads to finish before continuing.
      for thread in thread_list:
         thread.join()
      #
      return retcode
   #
#
